﻿namespace SpikingAnalyticsLib.ResponseFingerPrint

open System.IO
open System.Text
open SpikingNeuronLib
open SpikingAnalyticsLib
open SpikingAnalyticsLib.GeneralUtilities
open SpikingAnalyticsLib.IOUtilities

//=================================================================================================

/// A profile of the neural response for each neuron in the network
// The profile for one neuron is a response histogram
// The windowSize parameter is to support the OLD WindowMap algorithm
type FrameProfile private(windowSize, frameSize, numFrames, pattern:Pattern option, patternName, frameProfile) =

    // a unique file header to ensure a compressed file contains FrameProfile data
    static let uniqueID = [| 12uy; 90uy; 33uy; 40uy; 170uy; 117uy; 143uy; 76uy; 149uy; 76uy; 191uy; 24uy; 17uy; 57uy; 32uy; 112uy; |]

    /// Generate a map between selected neurons and their reponse histograms
    /// where the response histogram is a map of spike counts keyed by temporal offset from a specified frame
    static let GenerateFrameProfiles frameSize totalNeurons selectedNeurons firingEventList =

        let selectionMask = CreateFilter selectedNeurons totalNeurons

        firingEventList
        |> Seq.filter (fun (time, nindex, isBackground) -> selectionMask.[nindex])        // process only selected neurons
        |> Seq.map (fun (time, nindex, isBackground) -> time % frameSize, nindex)       // remap time into each frame of size 'frameSize'
        |> Seq.fold (fun (dict:Map<int,int list>) (time, nindex) ->         // a list of reponse times for each neuron
                let times = dict.[nindex]
                Map.add nindex (time :: times) dict
            ) (PreAllocateMap totalNeurons [])      // prellocate the map to ensure that all neurons have a slot (even when not present in the data)
        |> Map.map (fun nindex times ->                                     // turn each list of response times into a separate map of counts keyed by time
                times
                |> Seq.fold (fun (dict:Map<int,int>) time ->
                        if dict.ContainsKey(time) then
                            let count = dict.[time]
                            Map.add time (count + 1) dict
                        else
                            printfn "Warning: response time (%d) was out of range (%d time slots allocated)" time dict.Count
                            dict
                    ) (PreAllocateMap frameSize 0)
            )

    static let CreateProfiles (engine:SpikingAnalyticsLib.CrossbarNetwork) (pattern:Pattern option) patternName backgroundFiringRate windowSize frameSize runSeconds verbose plasticityEnabled =

        let numFrames = (1000 / frameSize) * runSeconds

        // profile BOTH excitatory and inhibitory neurons
        let selectedNeurons = [for i in 0..engine.TotalNeurons-1 -> i]

        let inputPattern =
            if pattern.IsSome then
                Some(Stimulus.Create(pattern.Value, frameSize) :> IStimulus)
            else
                None

        if verbose then
            printfn "Selecting windows..."
            printf "\tcollecting data..."

        engine.Run(runSeconds, inputPattern, backgroundFiringRate, plasticityEnabled)

        if verbose then printf "analysing %d events..." engine.FiringEventList.Count
        let frameProfile = GenerateFrameProfiles frameSize engine.TotalNeurons selectedNeurons engine.FiringEventList
        if verbose then printfn "done"

        numFrames, frameProfile

    // Generate a sequence that represents each response histogram
    // Select all histograms (key = None) or a single histogram (key = {int value})
    // The index is included at the start of each row
    let GetHistograms (index:int option) (key:int option) (histograms:Map<int, Map<int, int>>) =
        histograms
        |> Map.toSeq
        |> Seq.map (fun (nindex, histogram) ->
            if key.IsNone || key.IsSome && (nindex = key.Value) then
                let counts =
                    histogram
                    |> Map.toSeq
                    |> Seq.map (fun (time, count) -> count)
                    |> Seq.toList
                Some(index, nindex, counts)
            else
                None
        )
        |> Seq.choose id

    // Generate a sequence of strings that represent each response histogram
    // Select all histograms (key = None) or a single histogram (key = {int value})
    // The index is printed at the start of each row (if supplied)
    let ExtractHistograms (index:int option) (key:int option) (histograms:Map<int, Map<int, int>>) =
        GetHistograms index key histograms
        |> Seq.map (fun (index, nindex, counts) ->
                let str = System.String.Join(" ", counts)
                if index.IsSome then
                    sprintf "%d %d %s" index.Value nindex str
                else
                    sprintf "%d %s" nindex str
            )

//---------------------------------------------------------------------------------------------------------
// Constructors

    new (engine:SpikingAnalyticsLib.CrossbarNetwork, pattern:Pattern option, patternName, backgroundFiringRate, verbose) =

        let windowSize = 8      // to support the OLD WindowMap algorithm
        let frameSize = 250
        let runSeconds = 100

        let numFrames, frameProfile = CreateProfiles engine pattern patternName backgroundFiringRate windowSize frameSize runSeconds verbose false

        new FrameProfile(windowSize, frameSize, numFrames, pattern, patternName, frameProfile)

    new (frameSize, windowSize, runSeconds, engine:SpikingAnalyticsLib.CrossbarNetwork, pattern:Pattern option, patternName, backgroundFiringRate, verbose, ?plasticityEnabled) =

        assert (frameSize >= 40 && frameSize <= 1000)
        assert (windowSize > 0 && windowSize < 20)
        assert (runSeconds > 0)

        // enable/disable learning (disabled by default)
        let enablePlasticity = if plasticityEnabled.IsSome then plasticityEnabled.Value else false
        let numFrames, frameProfile = CreateProfiles engine pattern patternName backgroundFiringRate windowSize frameSize runSeconds verbose enablePlasticity

        new FrameProfile(windowSize, frameSize, numFrames, pattern, patternName, frameProfile)

//---------------------------------------------------------------------------------------------------------
// Properties

    member this.WindowSize = windowSize     // to support the OLD WindowMap algorithm

    member this.FrameSize = frameSize

    member this.NumberOfFrames = numFrames

    member this.PatternName = patternName

    member this.Pattern = pattern

    /// The response histogram for each neuron
    member this.ResponseHistograms = frameProfile

//---------------------------------------------------------------------------------------------------------
// Methods

    /// Retrieve selected histograms
    /// Select all histograms (key = None) or a selected histogram (key = {neuron index})
    /// The index is included at the start of each row (if supplied)
    member this.SelectHistograms(?key:int, ?index:int) =
        GetHistograms index key frameProfile
        |> Seq.toList

    /// Print all histograms (key = None) or a selected histogram (key = {neuron index}) to the console
    /// The index is printed at the start of each row (if supplied)
    member this.PrintHistograms(?key:int, ?index:int) =
        ExtractHistograms index key frameProfile
        |> Seq.iter (fun histogram -> printfn "%s" histogram)

    /// Save histograms to the specified stream in a simple tabular format that is easy for importing into other programs
    /// Save all histograms (key = None) or a selected histogram (key = {neuron index})
    /// The index is printed at the start of each row (if supplied)
    member this.SaveHistograms(sw:StreamWriter, ?key:int, ?index:int) =

        ExtractHistograms index key frameProfile
        |> Seq.iter (fun histogram -> sw.WriteLine(histogram))

    /// Save histograms to the specified file in a simple tabular format that is easy for importing into other programs
    /// Save all histograms (key = None) or a selected histogram (key = {neuron index})
    /// The index is printed at the start of each row (if supplied)
    member this.SaveHistograms(outputFilePath, ?key:int, ?index:int) =
        use sw = File.CreateText(outputFilePath)

        ExtractHistograms index key frameProfile
        |> Seq.iter (fun histogram -> sw.WriteLine(histogram))

    /// Save the frame profile to the specified stream
    member this.Save(sw:StreamWriter) =
        let header = sprintf "%d %d %d %f %f \"%s\"" windowSize frameSize numFrames 0.0 0.0 patternName
        sw.WriteLine(header)

        let patternString =
            if pattern.IsSome then
                pattern.Value.EventList
                |> Seq.map (fun evt -> evt.Time, evt.NeuronIndex)
                |> StringifyWithSeparator " "
            else
                "None"
        sw.WriteLine(patternString)

        frameProfile
        |> Map.iter (fun (neuronIndex:int) neuronProfile ->
                sw.WriteLine(neuronIndex)
                let profileString =
                    neuronProfile
                    |> Map.toSeq
                    |> Seq.map (fun (timeOffset, spikeCount) -> spikeCount)
                    |> StringifyWithSeparator " "
                sw.WriteLine(profileString)
            )

    /// Save the frame profile to the specified file
    member this.Save(outputFilePath) =
        use sw = File.CreateText(outputFilePath)
        this.Save(sw)

    /// Save the frame profile to the specified compressed binary file
    member this.SaveCompressedFrameProfile(outputFilePath) =
        // save to memory stream
        use saveStream = new MemoryStream()
        use sw = new StreamWriter(saveStream, Encoding.UTF8)
        this.Save(sw)
        sw.Close()  // ensure the writer is flushed to the memory stream

        // compress the data
        use zippedStream = ZipData (saveStream.ToArray())

        // add unique signature to the header
        let combined = Array.concat [ uniqueID; zippedStream.ToArray(); ]

        // save the compressed data to a file
        File.WriteAllBytes(outputFilePath, combined)

    /// Fabricate a new frame profile by transforming the response histogram data from this frame profile
    member this.Transform(patternName, pattern, transform:Map<int, Map<int, int>> -> Map<int, Map<int, int>>) =
        let transformedFrameProfile = transform frameProfile
        new FrameProfile(windowSize, frameSize, numFrames, pattern, patternName, transformedFrameProfile)

    /// Fabricate a new frame profile by keeping only the response histograms for the specified neurons
    /// All other response histograms are zeroed
    member this.Filter(patternName, pattern, retainedNeuronHistograms:Set<int>) =
        let ProfileFilter (histograms:Map<int, Map<int, int>>) =
            histograms
            |> Map.map (fun neuronIndex histogram ->
                    if retainedNeuronHistograms.Contains neuronIndex then
                        histogram
                    else
                        PreAllocateMap frameSize 0
                )
        this.Transform(patternName, pattern, ProfileFilter)

//---------------------------------------------------------------------------------------------------------
// Static Members

    /// Load the frame profile from the specified stream
    static member Load(sr:StreamReader, verbose) =
        let headerLine = sr.ReadLine()
        // Note that initialThreshold, frameFrequencyThreshold are now always zero (parameters moved to WindowMap)
        let windowSize, frameSize, numFrames, initialThreshold, frameFrequencyThreshold, patternName = IOUtilities.ParseStringToTuple<int * int * int * float * float * string>(headerLine)

        let pattern =
            let patternString = sr.ReadLine()
            if patternString.StartsWith("None") then
                None
            else
                let pairs = IOUtilities.ParseStringToList<int * int>(patternString)
                Some(new Pattern(pairs))

        let indexStack, profilesList =
            // read the rest of the file and tokenise it
            seq {
                while (sr.Peek() >= 0) do
                    let line = sr.ReadLine()
                    let tokens = IOUtilities.Tokenize(List.ofArray (line.ToCharArray()))
                    yield tokens
            }
            |> Seq.fold (fun (indexStack:System.Collections.Generic.Stack<int>, profilesList:int list list) tokens ->
                    // create a list of profiles where each profile is a list (includes neuronIndex at the head)
                    if tokens.Length = 1 then       // start of new neuron response histogram
                        let index =
                            let items = IOUtilities.ConstructList<int> tokens
                            items.Head
                        indexStack.Push(index)
                        indexStack, profilesList
                    else
                        // expecting the profile data
                        let spikeCounts = IOUtilities.ConstructTupleList<int> tokens
                        let neuronIndex = indexStack.Pop()
                        let neuronList = neuronIndex :: spikeCounts
                            
                        indexStack, neuronList :: profilesList
                ) (new System.Collections.Generic.Stack<int>(), [])

        let frameProfile =
            profilesList
            |> Seq.fold (fun (outer:Map<int, Map<int, int>>) patternList ->
                    let neuronIndex, profile =
                        match patternList with
                        | neuronIndex :: profile -> neuronIndex, profile
                        | [] -> 0, []
                    
                    let profileMap =
                        profile
                        |> Seq.mapi (fun index count -> index, count)
                        |> Map.ofSeq

                    Map.add neuronIndex profileMap outer
                ) Map.empty

        if verbose then
            printfn "done"

        new FrameProfile(windowSize, frameSize, numFrames, pattern, patternName, frameProfile)

    /// Load the frame profile from the specified file
    static member Load(inputFilePath:string) =
        FrameProfile.Load(inputFilePath, false)

    /// Load the frame profile from the specified file
    static member Load(inputFilePath, verbose) =
        if verbose then
            printf "Loading %s..." (Path.GetFileName(inputFilePath))

        use sr = File.OpenText(inputFilePath)
        FrameProfile.Load(sr, verbose)

    /// Load the selected profiles from the specified folder
    static member LoadMany(folder, fileNameBuilder:string -> string, patternNames, verbose) =

        let LoadProfile networkNumber patternName =
            let fileName = fileNameBuilder patternName
            FrameProfile.Load(Path.Combine(folder, fileName), verbose)

        patternNames
        |> Seq.map (fun patternName -> LoadProfile fileNameBuilder patternName)
        |> Seq.toArray

    /// Load the frame profile from the specified compressed binary file
    static member LoadCompressedFrameProfile(inputFilePath, verbose) =
        // read the binary file
        let zipped = File.ReadAllBytes(inputFilePath)

        // check the first 16 bytes match the unique ID
        let headerCorrect = Array.forall2 (fun a b -> a = b) uniqueID zipped.[..15]

        if headerCorrect then

            // uncompress the data
            use loadStream = UnzipData zipped.[16..]    // skip the first 16 bytes

            // load from memory stream
            loadStream.Seek(0L, SeekOrigin.Begin) |> ignore
            use sr = new StreamReader(loadStream, Encoding.UTF8)
            FrameProfile.Load(sr, verbose)

        else
            raise (System.ArgumentException("The specified file does not contain compressed frame profile data."))

    /// Generate a frame profile data structure from the specified firing events
    // Note this is a convenience function to support older APIs
    static member GenerateFrameProfileData(frameSize, totalNeurons, selectedNeurons, firingEventList) =
        GenerateFrameProfiles frameSize totalNeurons selectedNeurons firingEventList
